/*******************************************************************************
*
* Copyright (C) 2013 - 2015, Jeffery Jiang, <china_jeffery@163.com>.
*
* THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND,
* EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
*
* Expect bugs
*
* Please use and enjoy. Please let me know of any bugs/mods/improvements
* that you have found/implemented and I will fix/incorporate them into this
* file.
*
*******************************************************************************/

#include "WinPipe.h"
#include <tchar.h>
#include "CmnHdr.hpp"
#include "IPCManager.h"
#include "InternalTypes.h"

#define EXIT_KEY 1001

unsigned int __stdcall WorkThread(void *param);
unsigned int __stdcall VirtualMainThread(void *param);


WinPipe::WinPipe() : m_ICOP(0) {
    m_bListening = FALSE;
    m_PipeOLs = NULL;
    m_pWorkThreads = NULL;
    m_iWorkThreadNum = 0;
    m_iMaxConcurrent = kIPCMaxConcurrent;
    m_pfnListenCallback = NULL;
    m_dwBufferSize = kBufSize;
    ZeroMemory(m_szIPCName, sizeof(m_szIPCName));
}

WinPipe::~WinPipe() {

}

BOOL WinPipe::StartListen(LPCTSTR szIPCName, PFN_LISTENING_CALLBACK pfnCallback) {
    if(m_bListening)
        return FALSE;

    if(!IPCNameAvailable(TCHARToAnsi(szIPCName).c_str()))
        return FALSE;

	StringCchCopy(m_szIPCName, MAX_PATH, szIPCName);
    m_pfnListenCallback = pfnCallback;

    TCHAR szPipeName[MAX_PATH] = {0};
	StringCchPrintf(szPipeName, MAX_PATH, TEXT("\\\\.\\pipe\\%s"), szIPCName); 

    m_PipeOLs = new PipeOverlapped[m_iMaxConcurrent];

    for (int i = 0; i < m_iMaxConcurrent; i++) {
        HANDLE hPipe = CreateNamedPipe(szPipeName,
                                       PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED,
                                       PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE,
                                       PIPE_UNLIMITED_INSTANCES,
                                       kBufSize,
                                       kBufSize,
                                       2000,
                                       NULL);

        if(hPipe == INVALID_HANDLE_VALUE) {
            SAFE_DELETE_ARRAY(m_PipeOLs);
            return FALSE;
        }

        if(!m_ICOP.AssociateDevice(hPipe, i)) {
            SAFE_DELETE_ARRAY(m_PipeOLs);
            return FALSE;
        }


        m_PipeOLs[i].hPipe = hPipe;

        BOOL bRet = ConnectNamedPipe(hPipe, (LPOVERLAPPED)&m_PipeOLs[i].ol);

        if(bRet) {
            m_PipeOLs[i].iOpCode = OP_READ;
        }
        else {
            DWORD dwGLE = GetLastError();

            if(dwGLE == ERROR_IO_PENDING) {
                m_PipeOLs[i].iOpCode = OP_CONNECT;
            }
            else if(dwGLE == ERROR_PIPE_CONNECTED){
                m_PipeOLs[i].iOpCode = OP_READ;
            }
        }
    }

    ProcessFinder processFinder(TH32CS_SNAPPROCESS);
    PROCESSENTRY32 pe = {sizeof(pe)};
    chASSERT(processFinder.ProcessFind(GetCurrentProcessId(), &pe));

    IPCDetail detail;
    strcpy_s(detail.szIPCName, MAX_PATH, TCHARToAnsi(m_szIPCName).c_str());
    detail.dwProcessID = pe.th32ProcessID;
    strcpy_s(detail.szProcessName, MAX_PATH, TCHARToAnsi(pe.szExeFile).c_str());
    chASSERT(AddIPCDetail(detail));

    SYSTEM_INFO si;
    GetSystemInfo(&si);
    m_iWorkThreadNum = si.dwNumberOfProcessors * 2 + 2;
    m_pWorkThreads = new HANDLE[m_iWorkThreadNum];

    for(int i = 0; i < m_iWorkThreadNum; i++) {
        char szThreadName[MAX_PATH] = {0};
        StringCchPrintfA(szThreadName, MAX_PATH, "libipc-WorkThread%d", i);
        m_pWorkThreads[i] = CREATETHREAD(NULL, 0, WorkThread, this, 0, NULL, szThreadName);
    }

    m_hVirtualMainThread = CREATETHREAD(NULL, 0, VirtualMainThread, this, 0, NULL, "libipc-VirtualMainThread");

    m_bListening = TRUE;

    return TRUE;
}

void WinPipe::StopListen(void) {
    if(!m_bListening)
        return;

    for(int i = 0; i < m_iWorkThreadNum; i++) {
        m_ICOP.PostStatus(EXIT_KEY, 0, &m_PipeOLs[i].ol);
    }

    WaitForSingleObject(m_hVirtualMainThread, INFINITE);

    for(int i = 0; i < m_iMaxConcurrent; i++) {
        if(m_PipeOLs[i].hPipe) {
            FlushFileBuffers(m_PipeOLs[i].hPipe);
            DisconnectNamedPipe(m_PipeOLs[i].hPipe);
            CloseHandle(m_PipeOLs[i].hPipe);
            m_PipeOLs[i].hPipe = NULL;
        }
    }

    SAFE_DELETE_ARRAY(m_PipeOLs);
    SAFE_DELETE_ARRAY(m_pWorkThreads);
}


DWORD WinPipe::GetIPCName(TCHAR *szName, DWORD dwMaxSize) {
    if(!m_szIPCName)
        return 0;

    if(!szName)
        return _tcslen(m_szIPCName);

	StringCchCopy(szName, dwMaxSize, m_szIPCName);
    return _tcslen(m_szIPCName);
}

BOOL WinPipe::SetBufferSize(DWORD dwSize) {
    if(m_bListening)
        return FALSE;

    m_dwBufferSize = dwSize;
    return TRUE;
}

BOOL WinPipe::SetMaxConcurrent(int concurrent) {
    if(m_bListening)
        return FALSE;

    if(concurrent <= 0 || concurrent > PIPE_UNLIMITED_INSTANCES)
        return FALSE;

    m_iMaxConcurrent = concurrent;

    return TRUE;
}

BOOL WinPipe::Reconnect(PipeOverlapped *pPipeOL) {
    if(NULL == pPipeOL)
        return FALSE;

    if(NULL == pPipeOL->hPipe)
        return FALSE;

    FlushFileBuffers(pPipeOL->hPipe);

    if(!DisconnectNamedPipe(pPipeOL->hPipe))
        return FALSE;

    pPipeOL->iOpCode = OP_CONNECT;

    if((!ConnectNamedPipe(pPipeOL->hPipe, &(pPipeOL->ol))) && (GetLastError() == ERROR_IO_PENDING))
        return TRUE;

    return FALSE;
}

BOOL WinPipe::ReadPipe(PipeOverlapped *pPipeOL) {
    if(NULL == pPipeOL)
        return FALSE;

    if(NULL == pPipeOL->hPipe)
        return FALSE;

    BOOL bReadRequest  = FALSE;
    DWORD dwGLE = 0;
    DWORD dwBytesRead = 0;

    pPipeOL->iOpCode = OP_READ;

    if(!(bReadRequest = ReadFile(pPipeOL->hPipe, pPipeOL->pvBuff, pPipeOL->dwBuffMaxSize, &dwBytesRead, &pPipeOL->ol))) {
        dwGLE = GetLastError();

        if(dwGLE == ERROR_IO_PENDING || dwGLE == ERROR_PIPE_LISTENING )
            bReadRequest = TRUE;
    }

    return bReadRequest;
}



unsigned int __stdcall WorkThread(void *param) {
    WinPipe *pWinPipe = (WinPipe *)param;
    if(!pWinPipe)
        return 1;

    DWORD dwRead = 0;
    ULONG_PTR pKey = NULL;
    OVERLAPPED *pOL = NULL;
    PipeOverlapped *pPipeOL = NULL;

    while(true) {
        BOOL bRet = pWinPipe->m_ICOP.GetStatus(&pKey, &dwRead, &pOL);
        pPipeOL = CONTAINING_RECORD(pOL, PipeOverlapped, ol);

        if(!pOL)
            break;

        if(pKey == EXIT_KEY)
            break;

        switch(pPipeOL->iOpCode) {
            case OP_CONNECT: {
                    if(!pWinPipe->ReadPipe(pPipeOL))
                        chASSERT(pWinPipe->Reconnect(pPipeOL));

                    break;
                }

            case OP_READ: {
                    if(!bRet || dwRead == 0) { // Client Close.
                        chASSERT(pWinPipe->Reconnect(pPipeOL));
                        break;
                    }

                    pPipeOL->dwDataSize = pPipeOL->ol.InternalHigh;

                    if(pWinPipe->m_pfnListenCallback) {
                        MSGREQ_HEADER reqHeader;
                        memcpy(&reqHeader, pPipeOL->pvBuff, sizeof(MSGREQ_HEADER));
                        LPVOID pvAnswer = NULL;
                        DWORD dwAnswerMaxSize = pWinPipe->m_dwBufferSize;

                        if(reqHeader.bNeedAnswer) {
                            if(!(pvAnswer = malloc(dwAnswerMaxSize)))
                                dwAnswerMaxSize = 0;
                        }

                        DWORD dwAnswerSize = pWinPipe->m_pfnListenCallback(pWinPipe, (char *)pPipeOL->pvBuff + reqHeader.dwHeaderSize, reqHeader.dwContentSize, reqHeader.bNeedAnswer, pvAnswer, dwAnswerMaxSize);

						if(reqHeader.bNeedAnswer && dwAnswerSize > 0 && pvAnswer && dwAnswerMaxSize > 0) {
                            pPipeOL->iOpCode = OP_WRITE;

                            if(!WriteFile(pPipeOL->hPipe, pvAnswer, (dwAnswerSize <= dwAnswerMaxSize) ? dwAnswerSize : dwAnswerMaxSize, NULL, &pPipeOL->ol)) {
                                if(GetLastError() != ERROR_IO_PENDING) {
                                    SAFE_FREE(pvAnswer);
                                    break;
                                }
                            }
						}
                        else {
                            if(!pWinPipe->ReadPipe(pPipeOL))
                                chASSERT(pWinPipe->Reconnect(pPipeOL));
                        }

                        SAFE_FREE(pvAnswer);
                    }

                    break;
                }

            case OP_WRITE: {
                    if(!pWinPipe->ReadPipe(pPipeOL))
                        chASSERT(pWinPipe->Reconnect(pPipeOL));

                    break;
                }
        }
    }

    return 0;
}

unsigned int __stdcall VirtualMainThread(void *param) {
    WinPipe *pWinPipe = (WinPipe *)param;

    WaitForMultipleObjects(pWinPipe->m_iWorkThreadNum, pWinPipe->m_pWorkThreads, TRUE, INFINITE);

    return 1;
}
